#ifndef PCB_DETECTOR_H
#define PCB_DETECTOR_H

#include <iostream>
#include <vector>
#include <string>
#include <opencv2/opencv.hpp>

extern "C" {
#include <onnxruntime_c_api.h>
}

// 检测结果结构体
struct Detection {
    cv::Rect bbox;
    float confidence;
    int class_id;
    std::string class_name;
};

class PCBDetector {
private:
    // 配置参数
    float confidence_threshold_;
    float nms_threshold_;
    
    // 类别映射
    std::map<int, std::string> class_mapping_;
    
    // 颜色映射
    std::vector<cv::Scalar> colors_;
    
    // ONNX Runtime相关
    const OrtApi* ort_api_;
    OrtEnv* env_;
    OrtSession* session_;
    OrtAllocator* allocator_;
    
    // 输入输出信息
    char* input_name_;
    std::vector<char*> output_names_;
    size_t num_output_nodes_;

    // 私有方法声明
    void initOrt(const std::string& model_path);
    std::vector<Detection> postprocessResults(float* boxes_data, float* labels_data, float* scores_data, 
                                             int num_detections, int orig_width, int orig_height);
    cv::Mat drawDetections(const cv::Mat& image, const std::vector<Detection>& detections);
    void cleanup();
    bool checkStatus(OrtStatus* status);

public:
    // 构造函数
    PCBDetector(const std::string& model_path, 
                float confidence_threshold = 0.7f, 
                float nms_threshold = 0.3f);
    
    // 析构函数
    ~PCBDetector();
    
    // 检测方法
    std::vector<Detection> detectFromPath(const std::string& image_path, cv::Mat& result_image);
    std::vector<Detection> detect(const cv::Mat& input_image, cv::Mat& result_image);
};

#endif // PCB_DETECTOR_H